from unstructured.partition.pdf import partition_pdf
from pdf2image import convert_from_path
import os
from PIL import Image
import pytesseract
from typing import List, Dict

class PDFProcessor:
    def __init__(self, upload_dir: str):
        self.upload_dir = upload_dir
        os.makedirs(upload_dir, exist_ok=True)
    
    def process_pdf(self, pdf_path: str) -> List[Dict]:
        # 提取文本和表格
        elements = partition_pdf(pdf_path)
        
        # 提取图片
        images = convert_from_path(pdf_path)
        
        processed_data = []
        
        # 处理文本和表格
        for element in elements:
            if hasattr(element, "text"):
                processed_data.append({
                    "type": "text",
                    "content": element.text,
                    "page": getattr(element, "page_number", 1)
                })
        
        # 处理图片
        for idx, image in enumerate(images):
            # 使用OCR提取图片中的文本
            image_text = pytesseract.image_to_string(image, lang='chi_sim')
            
            processed_data.append({
                "type": "image",
                "content": image_text,
                "image": image,
                "page": idx + 1
            })
            
        return processed_data 